#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@File    ：es_client.py
@Author  ：平
@Date    ：2025/9/28 15:53 
"""

from elasticsearch import Elasticsearch

from app.config.config import settings
import logging

logger = logging.getLogger(__name__)

es_client = Elasticsearch(settings.ELASTICSEARCH_URL,
                          basic_auth=(settings.ELASTICSEARCH_USERNAME, settings.ELASTICSEARCH_PASSWORD))

# 初始化索引
# Elasticsearch 索引模板（含 settings + mappings）
# 用于存储“带文本、向量、元数据”的文档，例如 PDF、Office 等文件的解析结果
INDEX_MAPPING = {
    # ===== 索引级配置 =====
    "settings": {
        # 文本分析器：默认采用 ik_smart（粗粒度中文分词），并统一转小写
        "analysis": {
            "analyzer": {
                "default": {
                    "type": "ik_smart",
                    "filter": ["lowercase"]
                }
            }
        },
        "number_of_shards": 3,  # 主分片数，影响写入吞吐与并行度
        "number_of_replicas": 1,  # 副本数，保证高可用
    },
    # ===== 字段结构定义 =====
    "mappings": {
        "properties": {
            # 1. 元数据嵌套对象：存放文件级或页面级元信息
            "metadata": {
                "properties": {
                    "aigc": {  # 是否 AIGC 生成内容（yes/no 或其他标记）
                        "type": "text",
                        "fields": {
                            "keyword": {  # 多字段：保留不分词版本，用于精确聚合/过滤
                                "type": "keyword",
                                "ignore_above": 256  # 超过 256 字符直接忽略
                            }
                        }
                    },
                    "creationdate": {"type": "date"},  # 文件创建时间
                    "creator": {  # 作者
                        "type": "text",
                        "fields": {"keyword": {"type": "keyword", "ignore_above": 256}}
                    },
                    "moddate": {"type": "date"},  # 最后修改时间
                    "page": {"type": "long"},  # 当前页码（从 1 开始）
                    "page_label": {  # 页面标签（如罗马数字 i、ii）
                        "type": "text",
                        "fields": {"keyword": {"type": "keyword", "ignore_above": 256}}
                    },
                    "producer": {  # 生成工具（如 Acrobat 9.0）
                        "type": "text",
                        "fields": {"keyword": {"type": "keyword", "ignore_above": 256}}
                    },
                    "source": {  # 原始文件路径或 URL
                        "type": "text",
                        "fields": {"keyword": {"type": "keyword", "ignore_above": 256}}
                    },
                    "total_pages": {"type": "long"},  # 文件总页数
                    # 主键id
                    "id": {
                        "type": "long",
                        "fields": {"keyword": {"type": "keyword", "ignore_above": 256}}
                    },
                    # 标题
                    "title": {
                        "type": "text",
                        "fields": {"keyword": {"type": "keyword", "ignore_above": 256}}
                    },
                    # 描述
                    "description": {
                        "type": "text",
                        "fields": {"keyword": {"type": "keyword", "ignore_above": 256}}
                    },
                    # 作者
                    "author": {
                        "type": "text",
                        "fields": {"keyword": {"type": "keyword", "ignore_above": 256}}
                    },
                    #  出版社
                    "publisher": {
                        "type": "text",
                        "fields": {"keyword": {"type": "keyword", "ignore_above": 256}}
                    },
                    # 年份
                    "publishYear": {
                        "type": "long",
                        "fields": {"keyword": {"type": "keyword", "ignore_above": 256}}
                    },
                    # 分类
                    "category": {
                        "type": "text",
                        "fields": {"keyword": {"type": "keyword", "ignore_above": 256}}
                    },
                    # 标签
                    "tags": {
                        "type": "text",
                        "fields": {"keyword": {"type": "keyword", "ignore_above": 256}}
                    },
                    # 摘要
                    "summary": {
                        "type": "text",
                        "fields": {"keyword": {"type": "keyword", "ignore_above": 256}}
                    }
                }
            },

            # 2. 文本内容：PDF/Office 解析后的纯文本，支持全文检索
            "text": {
                "type": "text",
                "fields": {
                    "keyword": {  # 保留 256 字以内的精确值，用于排序/聚合
                        "type": "keyword",
                        "ignore_above": 256
                    }
                }
            },

            # 3. 稠密向量：用于语义搜索 / 向量召回
            "vector": {
                "type": "dense_vector",
                "dims": settings.EMBEDDING_MODEL_DIM,  # 维度由外部配置决定，如 768、1024
                "index": True,  # 启用近似最近邻索引
                "similarity": "cosine",  # 相似度算法：余弦
            }
        }
    }
}


# 初始化索引
def set_up():
    try:
        # 判断索引是否存在
        if not es_client.indices.exists(index=settings.ELASTICSEARCH_INDEX_NAME):
            logger.info(f"索引{settings.ELASTICSEARCH_INDEX_NAME}不存在，正在创建索引...")
            es_client.indices.create(index=settings.ELASTICSEARCH_INDEX_NAME, body=INDEX_MAPPING)
            logger.info(f"索引{settings.ELASTICSEARCH_INDEX_NAME}创建成功")
        else:
            logger.info(f"索引{settings.ELASTICSEARCH_INDEX_NAME}已存在")
    except Exception as e:
        logger.exception(f"索引{settings.ELASTICSEARCH_INDEX_NAME}创建失败:{e}")


# 创建索引
set_up()

if __name__ == '__main__':
    pass
    # response = es_client.search(index=settings.ELASTICSEARCH_INDEX_NAME, body={
    #     "query": {
    #         "match": {
    #             "text": "RAG"
    #         }
    #     }
    # })
    # for hit in response.get("hits", {}).get("hits", []):
    #     print(hit.get("_source", {}).get("text"))

    # response = es_client.search(
    #     index=settings.ELASTICSEARCH_INDEX_NAME,
    #     body={
    #         "query": {
    #            "term":{
    #                "metadata.id":1977571842822717442
    #            }
    #         }
    #     }
    # )
    # for hit in response.get("hits", {}).get("hits", []):
    #     print(hit.get("_source", {}).get("metadata").get('id'))
